# -*- coding: utf-8 -*-

# --- 导入新增的配置 ---
from agno.models.openai import OpenAILike

from configs.basic_config import ALI_GENERATION_MODEL, AGENT_URL, ALI_DASHSCOPE_API_KEY
import os

# --- 修正导入路径 (假设 agno.os 是正确的) ---
# 如果上面的导入不起作用，请尝试下面的替代方案之一：
# from agno.os import AgentOS # 假设这是正确的导入方式
# from agno.os.agentos import AgentOS # 如果 AgentOS 在 agentos.py 中
# 请根据实际的 agno 库结构调整导入

from agno.agent.agent import Agent
from agno.db.sqlite import SqliteDb

from agno.os import AgentOS
from agno.tools.notion import NotionTools
from agno.workflow.step import Step, StepInput, StepOutput
from agno.workflow.workflow import Workflow
from pydantic import BaseModel


# Pydantic model for classification output
class ClassificationResult(BaseModel):
    query: str
    tag: str
    message: str


# Agents
# --- 修改模型参数 ---
notion_agent = Agent(
    name="Notion 管理员",
    model=OpenAILike(id=ALI_GENERATION_MODEL, name=ALI_GENERATION_MODEL, api_key=ALI_DASHSCOPE_API_KEY, base_url=AGENT_URL),
    tools=[
        NotionTools(
            api_key=os.getenv("NOTION_API_KEY", ""),
            database_id=os.getenv("NOTION_DATABASE_ID", ""),
        )
    ],
    instructions=[
        "你是一个 Notion 页面管理员。",
        "你将收到带有查询和预分类标签的指令。",
        "关键：仅使用指令中提供的精确标签。不要创建新标签或修改标签名称。",
        "有效标签是：travel, tech, general-blogs, fashion, documents",
        "工作流程：",
        "1. 搜索具有所提供精确标签的现有页面",
        "2. 如果存在页面：使用新查询内容更新该页面",
        "3. 如果不存在页面：使用提供的精确标签创建新页面",
        "始终保留指令中给出的确切标签名称。",
    ],
)


# Executor functions
# Step 1: 自定义分类函数，将标签分配给查询
def classify_query(step_input: StepInput) -> StepOutput:
    """
    将用户查询分类到预定义的标签之一。

    可用标签：travel, tech, general-blogs, fashion, documents
    """
    # 从 step_input 获取用户查询
    query = step_input.input

    # 创建一个代理来分类查询
    # --- 修改模型参数 ---
    classifier_agent = Agent(
        model=OpenAILike(id=ALI_GENERATION_MODEL, name=ALI_GENERATION_MODEL, api_key=ALI_DASHSCOPE_API_KEY, base_url=AGENT_URL),
        instructions=[
            "你是一个查询分类器。",
            "将给定的查询分类到以下标签之一：travel, tech, general-blogs, fashion, documents",
            "只回复标签名称，别无其他。",
            "分类规则：",
            "- travel: 与目的地、旅游、旅行、地点、酒店、旅行指南、游览地相关的内容",
            "- tech: 编程、软件、人工智能、机器学习、编码、开发、技术主题",
            "- fashion: 服装、时尚、潮流、搭配、时尚产业",
            "- documents: 简历、CV、报告、官方文件、合同",
            "- general-blogs: 个人想法、观点、生活建议、杂项内容",
            "",
            "示例：",
            "- '意大利最佳游览地' -> travel",
            "- '越南河江环线旅游指南' -> travel",
            "- '添加旅游指南网站链接' -> travel",
            "- '如何构建 React 应用' -> tech",
            "- '人工智能和机器学习的兴起' -> tech",
            "- '2025年时尚趋势' -> fashion",
            "- '我的简历和 CV' -> documents",
            "- '关于生活的随机想法' -> general-blogs",
        ],
    )

    # 获取分类结果
    response = classifier_agent.run(query)
    tag = response.content.strip().lower()

    # 验证标签
    valid_tags = ["travel", "tech", "general-blogs", "fashion", "documents"]
    if tag not in valid_tags:
        tag = "general-blogs"  # 默认回退

    # 使用 Pydantic 模型返回结构化数据
    result = ClassificationResult(
        query=str(query), tag=tag, message=f"查询已分类为：{tag}"
    )

    return StepOutput(content=result)


# 自定义函数，准备传递给 Notion 代理的输入
def prepare_notion_input(step_input: StepInput) -> StepOutput:
    """
    提取分类结果并将其格式化以供 Notion 代理使用。
    """
    # 从上一步（分类查询）获取分类结果
    previous_output = step_input.previous_step_content

    # 如果是字典，则解析为我们的 Pydantic 模型
    if isinstance(previous_output, dict):
        classification = ClassificationResult(**previous_output)
    elif isinstance(previous_output, str):
        # 如果是字符串，尝试解析它或使用原始输入
        import json

        try:
            classification = ClassificationResult(**json.loads(previous_output))
        except (json.JSONDecodeError, TypeError, KeyError, ValueError):
            classification = ClassificationResult(
                query=str(step_input.input),
                tag="general-blogs",
                message="无法解析分类",
            )
    else:
        classification = previous_output

    # 为 Notion 代理创建清晰的指令，明确要求使用标签
    instruction = f"""处理此已分类的查询：

        查询: {classification.query}
        标签: {classification.tag}

        重要：你必须使用标签 "{classification.tag}" (有效标签之一: travel, tech, general-blogs, fashion, documents)。
        不要创建新标签。必须精确使用 "{classification.tag}"。

        指令：
        1. 使用 search_pages 工具查找标签为 "{classification.tag}" 的页面
        2. 如果页面存在：使用 update_page 添加查询内容
        3. 如果没有页面存在：使用 create_page 创建，标题为 "My {classification.tag.title()} Collection"，标签为 "{classification.tag}"，并将查询作为内容

        标签必须完全一致：{classification.tag}
    """

    return StepOutput(content=instruction)


# Steps
classify_step = Step(
    name="分类查询",
    executor=classify_query,
    description="将用户查询分类到标签类别",
)

notion_prep_step = Step(
    name="准备 Notion 输入",
    executor=prepare_notion_input,
    description="将分类结果格式化为 Notion 代理的输入",
)

notion_step = Step(
    name="管理 Notion 页面",
    agent=notion_agent,
    description="根据查询和标签创建或更新 Notion 页面",
)

# 创建工作流
query_to_notion_workflow = Workflow(
    name="query-to-notion-workflow",
    description="分类用户查询并在 Notion 中组织它们",
    # db=SqliteDb(
    #     session_table="workflow_session",
    #     db_file="tmp/workflow.db",
    # ),
    steps=[classify_step, notion_prep_step, notion_step],
)

# 初始化 AgentOS
agent_os = AgentOS(
    description="查询分类和 Notion 组织系统",
    workflows=[query_to_notion_workflow],
)
app = agent_os.get_app()

if __name__ == "__main__":
    # --- 注意：如果上面的导入有问题，请确认实际的导入方式 ---
    agent_os.serve(app="notion_manager:app", reload=True) # 这里可能需要根据实际的 app 名称调整